/**
 * Copyright 2020 Inc.
 **/
package com.myz.cloud.rocketmq.producer.mq.transaction;

import lombok.extern.slf4j.Slf4j;
import org.apache.rocketmq.spring.annotation.RocketMQTransactionListener;
import org.apache.rocketmq.spring.core.RocketMQLocalTransactionListener;
import org.apache.rocketmq.spring.core.RocketMQLocalTransactionState;
import org.apache.rocketmq.spring.support.RocketMQHeaders;
import org.springframework.messaging.Message;

import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * @author maoyz0621 on 2020/8/12
 * @version v1.0
 */
@RocketMQTransactionListener(txProducerGroup = "${rocketmq.transaction.group}")
@Slf4j
public class MqConsumerTransaction implements RocketMQLocalTransactionListener {

    private AtomicInteger transactionIndex = new AtomicInteger(0);

    private ConcurrentHashMap<String, Integer> cache = new ConcurrentHashMap<>();

    @Override
    public RocketMQLocalTransactionState executeLocalTransaction(Message msg, Object arg) {
        String id = (String) msg.getHeaders().get(RocketMQHeaders.TRANSACTION_ID);
        int val = transactionIndex.getAndIncrement();
        int status = val % 3;
        cache.put(id, status);
        log.info("executeLocalTransaction id = {}, tag = {}, msg = {}, arg = {}", msg.getHeaders().get("rocketmq_TRANSACTION_ID"), msg.getHeaders().get("rocketmq_TAGS"),msg.getPayload().toString(), arg);
        if (status == 0) {
            return RocketMQLocalTransactionState.COMMIT;  // 提交事务，它允许消费者消费此消息
        } else if (status == 1) {
            return RocketMQLocalTransactionState.ROLLBACK;  // 回滚事务，它代表该消息将被删除，不允许被消费
        }
        return RocketMQLocalTransactionState.UNKNOWN;  // 中间状态，它代表需要检查消息队列来确定状态
    }

    @Override
    public RocketMQLocalTransactionState checkLocalTransaction(Message msg) {
        String id = (String) msg.getHeaders().get(RocketMQHeaders.TRANSACTION_ID);
        RocketMQLocalTransactionState retState = RocketMQLocalTransactionState.COMMIT;
        Integer status = cache.get(id);
        if (status != null) {
            switch (status) {
                case 0:
                    retState = RocketMQLocalTransactionState.UNKNOWN;
                    break;
                case 1:
                    retState = RocketMQLocalTransactionState.COMMIT;
                    break;
                case 2:
                    retState = RocketMQLocalTransactionState.ROLLBACK;
                    break;
                default:
                    retState = RocketMQLocalTransactionState.UNKNOWN;
            }
        }
        log.info("checkLocalTransaction id = {}, tag = {}", msg.getHeaders().get("rocketmq_TRANSACTION_ID"), msg.getHeaders().get("rocketmq_TAGS"));

        return retState;
    }
}